#include "b_tree.h"
#include "buffer_pool.h"

#include <stdio.h>

void b_tree_init(const char *filename, BufferPool *pool) {
    init_buffer_pool(filename, pool);
     
    if(pool->file.length != 0)
        return;
    
    BCtrlBlock* ctl = (BCtrlBlock*)get_page(pool, 0);
    ctl->root_node = -1;
    ctl->free_node_head = -1;
    release(pool, 0);
}

void b_tree_close(BufferPool *pool) {
    close_buffer_pool(pool);
}

RID b_tree_search(BufferPool *pool, void *key, size_t size, b_tree_ptr_row_cmp_t cmp) {
    RID rid; // default empty
    get_rid_block_addr(rid) = -1; 
    get_rid_idx(rid) = 0;
    
    BCtrlBlock* ctl = (BCtrlBlock*)get_page(pool, 0); // get ctl
    
    int i = 0;
    off_t child = ctl->root_node, next_child = -1;
    while(child != -1){
        BNode *node = (BNode*)get_page(pool, child);
        if(node->leaf == 1){ // is leaf
            for(i = 0; i < node->n; i++)
                if(cmp(key, size, node->row_ptr[i]) == 0){ // find
                    rid = node->row_ptr[i];
                    break;
                }
            release(pool, child);
            child = -1;
        }else{ // not a leaf
            next_child = node->child[node->n]; // default the biggest
            for(i = 0; i < node->n; i++)
                if(cmp(key, size, node->row_ptr[i]) < 0){ // find the first bigger one
                    next_child = node->child[i];
                    break;
                }
            release(pool, child);
            child = next_child;
        }
    }

    release(pool, 0); // release ctl

    return rid;
}

NewChildEntry b_tree_insert_recursively(BufferPool* pool, RID rid, b_tree_row_row_cmp_t cmp, b_tree_insert_nonleaf_handler_t insert_handler, off_t addr) {
    /**
     * get insert_pos and use b_tree_insert_recursively and compose new_child_entry
     * b_tree_insert_recursively ensure that current node is modified, and return NewChildEntry if split:
     * insert new_child_entry from below node into current node, consider split or not, and create rst to return
     */
    
    NewChildEntry rst; // default empty
    get_rid_block_addr(rst.rid) = -1;
    get_rid_idx(rst.rid) = 0;

    BNode* node = (BNode*)get_page(pool, addr);

    /**
     * recursively insert and get new_child_entry
     */
    // get insert_pos
    int insert_pos = 0;
    for (insert_pos = 0; insert_pos < node->n; insert_pos++)
        if (cmp(rid, node->row_ptr[insert_pos]) < 0) 
            break;
    // recursively insert
    NewChildEntry new_child_entry;
    if (node->leaf == 0) {
        off_t next_addr = node->child[insert_pos]; // get next_addr
        release(pool, addr); // !!! buffer pool may be not enough, so must release first
        new_child_entry = b_tree_insert_recursively(pool, rid, cmp, insert_handler, next_addr); // recursively, insert entry
        if (get_rid_block_addr(new_child_entry.rid) == -1) // if do not split, directly return
            return rst;
        node = (BNode*)get_page(pool, addr); // reget node
    }else{ // leaf
        new_child_entry.rid = rid;
        new_child_entry.addr = -1;
    }

    /**
     * if split below, insert new_child_entry into node
     */
    int i = 0;
    if(node->n < 2 * DEGREE){ // has space, no more split
        for(i = node->n; i > insert_pos; i--){
            node->row_ptr[i] = node->row_ptr[i - 1];
            node->child[i + 1] = node->child[i];
        }
        node->row_ptr[insert_pos] = new_child_entry.rid;
        node->child[insert_pos + 1] = new_child_entry.addr;
        node->n++;
    }else{ // split
        // get array_to_split after insert
        RID row_ptr_to_split[2 * DEGREE + 1];
        off_t child_to_split[2 * DEGREE + 2];
        for(i = 0; i < insert_pos; i++){
            row_ptr_to_split[i] = node->row_ptr[i];
            child_to_split[i] = node->child[i];
        }
        child_to_split[insert_pos] = node->child[insert_pos];
        row_ptr_to_split[insert_pos] = new_child_entry.rid;
        child_to_split[insert_pos + 1] = new_child_entry.addr;
        for(i = insert_pos + 1; i < 2 * DEGREE + 1; i++){
            row_ptr_to_split[i] = node->row_ptr[i - 1];
            child_to_split[i + 1] = node->child[i];
        }
        // ask for space for new node
        BCtrlBlock* ctl = (BCtrlBlock*)get_page(pool, 0);
        off_t new_node_addr = -1;
        BNode* new_node = NULL;
        if(ctl->free_node_head != -1) {
            new_node_addr = ctl->free_node_head;
            new_node = (BNode*)get_page(pool, new_node_addr);
            ctl->free_node_head = new_node->next;
        }else{
            new_node_addr = pool->file.length;
            new_node = (BNode*)get_page(pool, new_node_addr);
        }
        release(pool, 0);
        // split into node
        node->n = DEGREE;
        for(i = 0; i < node->n; i++){
            node->row_ptr[i] = row_ptr_to_split[i];
            node->child[i] = child_to_split[i];
        }
        node->child[node->n] = child_to_split[node->n];
        // split into new_node.
        if(node->leaf){ // leaf=>copy up
            new_node->n = DEGREE + 1;
            for(i = 0; i < new_node->n; i++){
                new_node->row_ptr[i] = row_ptr_to_split[DEGREE + i];
                new_node->child[i] = child_to_split[DEGREE + i];
            }
            new_node->child[new_node->n] = child_to_split[DEGREE + new_node->n];
            new_node->leaf = 1;
            new_node->next = node->next;
            node->next = new_node_addr;
        }else{ // non-leaf=>push up
            new_node->n = DEGREE;
            for(i = 0; i < new_node->n; i++){
                new_node->row_ptr[i] = row_ptr_to_split[DEGREE + 1 + i];
                new_node->child[i] = child_to_split[DEGREE + 1 + i];
            }
            new_node->child[new_node->n] = child_to_split[DEGREE + 1 + new_node->n];
            new_node->leaf = 0;
        }
        release(pool, new_node_addr);
        // compose rst
        if(node->leaf == 0) // not a leaf
            rst.rid = row_ptr_to_split[DEGREE];
        else // leaf
            rst.rid = insert_handler(row_ptr_to_split[DEGREE]);
        rst.addr = new_node_addr;
    }

    release(pool, addr);
    return rst;
}

RID b_tree_insert(BufferPool* pool, RID rid, b_tree_row_row_cmp_t cmp, b_tree_insert_nonleaf_handler_t insert_handler) {
    /**
     * get root, if failed, create one
     * use b_tree_insert_recursively
     * if generate new root, update ctl
     */

    BCtrlBlock* ctl = (BCtrlBlock*)get_page(pool, 0);

    /**
     * get root_addr, if failed, create a new root
     */
    off_t root_addr = ctl->root_node;
    if (root_addr == -1) { // empty
        BNode* root = NULL;
        // get root and update ctl->free_node_head
        if(ctl->free_node_head != -1) {
            root_addr = ctl->free_node_head;
            root = (BNode*)get_page(pool, root_addr);
            ctl->free_node_head = root->next;
        }else{
            root_addr = pool->file.length;
            root = (BNode*)get_page(pool, root_addr);
        }
        // update root & ctl->root_node
        root->n = 1;
        root->next = -1;
        root->child[0] = -1;
        root->child[1] = -1;
        root->row_ptr[0] = rid;
        root->leaf = 1;
        release(pool, root_addr);
        ctl->root_node = root_addr;

        release(pool, 0);
        return rid;
    }

    /**
     * insert recursively
     */
    release(pool, 0); // release first, in case buffer_pool is not enough
    NewChildEntry new_child_entry = b_tree_insert_recursively(pool, rid, cmp, insert_handler, root_addr);
    ctl = (BCtrlBlock*)get_page(pool, 0);

    /**
     * check if need to update root
     */
    if(get_rid_block_addr(new_child_entry.rid) == -1){ // no new root
        release(pool, 0);
        return rid;
    }
    // get root and update ctl->free_node_head
    BNode* new_root;
    off_t new_root_addr;
    if(ctl->free_node_head != -1) {
        new_root_addr = ctl->free_node_head;
        new_root = (BNode*)get_page(pool, new_root_addr);
        ctl->free_node_head = new_root->next;
    }else{
        new_root_addr = pool->file.length;
        new_root = (BNode*)get_page(pool, new_root_addr);
    }
    // update root & ctl->root_node
    new_root->n = 1;
    new_root->next = -1;
    new_root->child[0] = ctl->root_node;
    new_root->child[1] = new_child_entry.addr;
    new_root->row_ptr[0] = new_child_entry.rid;
    new_root->leaf = 0;
    release(pool, new_root_addr);
    ctl->root_node = new_root_addr;

    release(pool, 0);
    return rid;
}

void b_tree_delete_recursively(BufferPool* pool, RID rid, b_tree_row_row_cmp_t cmp, b_tree_insert_nonleaf_handler_t insert_handler, b_tree_delete_nonleaf_handler_t delete_handler, off_t addr, off_t left_sibling_addr, off_t right_sibling_addr, off_t parent_addr, int delete_pos) {
    /**
     * get next_delete_info and use b_tree_delete_recursively
     * b_tree_delete_recursively ensure that current node is modified, as well as delete_pos in its parent:
     * root->directly return, else->consider delete/redistribution/merge
     */
    
    BNode *node = (BNode*)get_page(pool, addr);

    /**
     * recursively delete
     */
    int next_delete_pos = 0;
    if(node->leaf){ // leaf
        // get next_delete_pos
        for(next_delete_pos = 0; next_delete_pos < node->n; next_delete_pos++)
            if(cmp(rid, node->row_ptr[next_delete_pos]) <= 0)
                break;
        if(cmp(rid, node->row_ptr[next_delete_pos]) != 0){ // not find
            release(pool, addr);
            return;
        }
        // delete rid
        int i = 0;
        for(i = next_delete_pos; i < node->n - 1; i++){
            node->child[i + 1] = node->child[i + 2];
            node->row_ptr[i] = node->row_ptr[i + 1];
        }
        node->n--;
    }else{ // non-leaf
        // get next_delete_pos
        for(next_delete_pos = 0; next_delete_pos < node->n; next_delete_pos++)
            if(cmp(rid, node->row_ptr[next_delete_pos]) < 0)
                break;
        // get next_addr, next_left_sibling_addr, next_right_sibling_addr
        off_t next_addr = node->child[next_delete_pos], 
            next_left_sibling_addr = next_delete_pos == 0 ? -1 : node->child[next_delete_pos - 1], 
            next_right_sibling_addr = next_delete_pos == node->n ? -1 : node->child[next_delete_pos + 1];
        // recursively delete
        release(pool, addr); // !!! buffer pool may be not enough, so must release first
        b_tree_delete_recursively(pool, rid, cmp, insert_handler, delete_handler, next_addr, next_left_sibling_addr, next_right_sibling_addr, addr, next_delete_pos);
        node = (BNode*)get_page(pool, addr);
    }

    /**
     * root, b_tree_delete will handle it
     */
    if(parent_addr == -1){
        release(pool, addr);
        return;
    }

    /**
     * if node->n is not enough, 
     * redistribute or merge, 
     * and adjust parent_node
     */
    if(node->n >= DEGREE){ // directly return
        release(pool, addr);
        return;
    }
    // get parent and siblings
    BNode *parent_node = parent_addr == -1 ? NULL : (BNode*)get_page(pool, parent_addr),
          *left_sibling = left_sibling_addr == -1 ? NULL : (BNode*)get_page(pool, left_sibling_addr),
          *right_sibling = right_sibling_addr == -1 ? NULL : (BNode*)get_page(pool, right_sibling_addr);
    // redistribute
    int j = 0;
    if(left_sibling != NULL && left_sibling->n - 1 >= DEGREE){ // left
        if(node->leaf){ // leaf => move and copy up
            // move
            for(j = node->n; j > 0; j--)
                node->row_ptr[j] = node->row_ptr[j - 1];
            // redistribute
            node->row_ptr[0] = left_sibling->row_ptr[left_sibling->n - 1];
            delete_handler(parent_node->row_ptr[delete_pos - 1]);
            parent_node->row_ptr[delete_pos - 1] = insert_handler(node->row_ptr[0]);
        }else{ // non-leaf => pull down and push up
            // move
            for(j = node->n; j > 0; j--){
                node->child[j + 1] = node->child[j];
                node->row_ptr[j] = node->row_ptr[j - 1];
            }
            node->child[1] = node->child[0];
            // redistribute
            node->child[0] = left_sibling->child[left_sibling->n];
            node->row_ptr[0] = parent_node->row_ptr[delete_pos - 1];
            parent_node->row_ptr[delete_pos - 1] = left_sibling->row_ptr[left_sibling->n - 1];
        }
        // update num
        left_sibling->n--;
        node->n++;
    }else if(right_sibling != NULL && right_sibling->n - 1 >= DEGREE){ // right
        if(node->leaf){ // leaf => move and copy up
            // redistribute
            node->row_ptr[node->n] = right_sibling->row_ptr[0];
            delete_handler(parent_node->row_ptr[delete_pos]);
            parent_node->row_ptr[delete_pos] = insert_handler(right_sibling->row_ptr[1]);
            // move
            for(j = 0; j < right_sibling->n - 1; j++)
                right_sibling->row_ptr[j] = right_sibling->row_ptr[j + 1];
        }else{ // non-leaf => pull down and push up
            // redistribute
            node->child[node->n + 1] = right_sibling->child[0];
            node->row_ptr[node->n] = parent_node->row_ptr[delete_pos];
            parent_node->row_ptr[delete_pos] = right_sibling->row_ptr[0];
            // move
            for(j = 0; j < right_sibling->n - 1; j++){
                right_sibling->child[j] = right_sibling->child[j + 1];
                right_sibling->row_ptr[j] = right_sibling->row_ptr[j + 1];
            }
            right_sibling->child[right_sibling->n - 1] = right_sibling->child[right_sibling->n];
        }
        // update num
        right_sibling->n--;
        node->n++;
    }
    // merge
    else if(left_sibling != NULL && left_sibling->n - 1 < DEGREE){ // left
        if(node->leaf){ // leaf => move and delete
            // move child
            for(j = 0; j < node->n; j++)
                left_sibling->row_ptr[left_sibling->n + j] = node->row_ptr[j];
            // delete
            delete_handler(parent_node->row_ptr[delete_pos - 1]);
        }else{ // non-leaf => pull down and move
            // pull down
            left_sibling->row_ptr[left_sibling->n] = parent_node->row_ptr[delete_pos - 1];
            // move child
            for(j = 0; j < node->n; j++){
                left_sibling->child[left_sibling->n + 1 + j] = node->child[j];
                left_sibling->row_ptr[left_sibling->n + 1 + j] = node->row_ptr[j];
            }
            left_sibling->child[left_sibling->n + 1 + node->n] = node->child[node->n];
        }
        // move parent
        for(j = delete_pos - 1; j < parent_node->n - 1; j++){
            parent_node->child[j + 1] = parent_node->child[j + 2];
            parent_node->row_ptr[j] = parent_node->row_ptr[j + 1];
        }
        // update num
        left_sibling->n = left_sibling->n + (!node->leaf) + node->n;
        parent_node->n--;
        // update ctl
        BCtrlBlock *ctl = (BCtrlBlock*)get_page(pool, 0);
        node->next = ctl->free_node_head;
        ctl->free_node_head = addr;
        release(pool, 0);
    }else if(right_sibling != NULL && right_sibling->n - 1 < DEGREE){ // right
        if(node->leaf){ // leaf => move and delete
            // move child
            for(j = 0; j < right_sibling->n; j++)
                node->row_ptr[node->n + j] = right_sibling->row_ptr[j];
            // delete
            delete_handler(parent_node->row_ptr[delete_pos]);
        }else{ // non-leaf => pull down and move
            // pull down
            node->row_ptr[node->n] = parent_node->row_ptr[delete_pos];
            // move child
            for(j = 0; j < right_sibling->n; j++){
                node->child[node->n + 1 + j] = right_sibling->child[j];
                node->row_ptr[node->n + 1 + j] = right_sibling->row_ptr[j];
            }
            node->child[node->n + 1 + right_sibling->n] = right_sibling->child[right_sibling->n];
        }
        // move parent
        for(j = delete_pos; j < parent_node->n - 1; j++){
            parent_node->child[j + 1] = parent_node->child[j + 2];
            parent_node->row_ptr[j] = parent_node->row_ptr[j + 1];
        }
        // update num
        node->n = node->n + (!node->leaf) + right_sibling->n;
        parent_node->n--;
        // update ctl
        BCtrlBlock *ctl = (BCtrlBlock*)get_page(pool, 0);
        right_sibling->next = ctl->free_node_head;
        ctl->free_node_head = right_sibling_addr;
        release(pool, 0);
    }
    // release parent and siblings
    if(parent_addr != -1) release(pool, parent_addr);
    if(left_sibling_addr != -1) release(pool, left_sibling_addr);
    if(right_sibling_addr != -1) release(pool, right_sibling_addr);

    release(pool, addr);
    return;
}

void b_tree_delete(BufferPool* pool, RID rid, b_tree_row_row_cmp_t cmp, b_tree_insert_nonleaf_handler_t insert_handler, b_tree_delete_nonleaf_handler_t delete_handler) {
    /**
     * get root, if failed, return
     * use b_tree_delete_recursively
     * if old root is empty, update ctl
     */
    
    BCtrlBlock *ctl = (BCtrlBlock*)get_page(pool, 0);

    /**
     * get root_addr, if empty, return directly
     */
    off_t root_addr = ctl->root_node;
    if(root_addr == -1){ // empty
        release(pool, 0);
        return;
    }

    /**
     * delete recursively
     */
    release(pool, 0); // release first, in case buffer_pool is not enough
    b_tree_delete_recursively(pool, rid, cmp, insert_handler, delete_handler, root_addr, -1, -1, -1, 0);
    ctl = (BCtrlBlock*)get_page(pool, 0);

    /**
     * check if need to update root
     */
    BNode *root = (BNode*)get_page(pool, root_addr);
    if(root->n == 0){ // empty
        root->next = ctl->free_node_head;
        ctl->free_node_head = root_addr;
        ctl->root_node = root->child[0];
    }
    release(pool, root_addr);

    release(pool, 0);
    return;
}